import torch
import gradio as gr

title = "基于YOLOV5的目标检测"
desc = "本项目基于yolov5进行实现检测功能，基于Gradio作为展示功能"


base_conf = 0.25
base_iou = 0.45

model = torch.hub.load("./","custom",path="runs/train/exp3/weights/best.pt",source="local")


def image_dect(image,conf,iou):
    model.conf = conf
    model.iou = iou
    results = model(image)
    return results.render()[0]

examples = [
    ['testData/1.jpg',0.65,0.7],
    ['testData/2.jpg',0.65,0.7]
        ]


gr.Interface(
    fn=image_dect,
    inputs = ["image",gr.Slider(minimum=0,maximum=1,value=base_conf),gr.Slider(minimum=0,maximum=1,value=base_iou)],
    outputs = ["image"],
    title = title,
    live=True,
    description = desc,
    examples=examples
).launch()